'''
Nelder-Mead方法
Author：Hanpeng Lin
Time: 2024/12/5 10:25
Email: linhp7@mail2.sysu.edu.cn
'''

import numpy as np
import matplotlib.pyplot as plt

def nelder_mead(func, x0, r=1, e=0.5, c=0.5, s=0.5, max_iter=1, tolerance=1e-100):
    """
    Nelder-Mead 算法实现。

    参数:
        func: 目标函数
        x0: 初始点数组，shape=(n+1, n)，n为变量维度
        r: 反射系数 (默认: 1)
        e: 扩展系数 (默认: 0.5)
        c: 收缩系数 (默认: 0.5)
        s: 压缩系数 (默认: 0.5)
        max_iter: 最大迭代次数 (默认: 100000)
        tolerance: 收敛精度 (默认: 1e-10)

    返回:
        points: 包含迭代结果的数组，格式为 [x, y, f]
        history_f: 每次迭代的目标函数值
    """
    X = np.array(x0)
    history_f = []  # 存储每次迭代的目标函数值

    for iter in range(max_iter):
        # 计算函数值并排序
        f_values = np.apply_along_axis(func, 1, X)
        order = np.argsort(f_values)
        X = X[order]

        # 记录当前最小点对应的目标函数值
        history_f.append(f_values[0])

        # 收敛判定：如果目标函数值变化很小
        if np.abs(f_values[0] - f_values[-1]) < tolerance:
            break

        # 计算质心
        centroid = np.mean(X[:-1], axis=0)

        # 反射
        xr = centroid + r * (centroid - X[-1])
        fr = func(xr)

        if fr < f_values[0]:
            # 扩展
            xe = centroid + e * (xr - centroid)
            fe = func(xe)

            if fe < f_values[0]:
                X[-1] = xe
            else:
                X[-1] = xr
        elif fr < f_values[-2]:
            X[-1] = xr
        else:
            if fr < f_values[-1]:
                X[-1] = xr

            # 收缩
            xs = centroid + c * (X[-1] - centroid)
            fs = func(xs)

            if fs < f_values[-1]:
                X[-1] = xs
            else:
                # 压缩
                X[1:] = X[0] + s * (X[1:] - X[0])

    # 最终点的坐标及函数值
    points = [X, func(X[0])]
    return points, history_f

# # 绘制收敛曲线
# def plot_convergence(history_f):
#     plt.figure(figsize=(10, 6))
#     plt.plot(range(len(history_f)), history_f, label='func(X[0])')
#     plt.xlabel('Iteration')
#     plt.ylabel('Function Value')
#     plt.title('Convergence Curve of func(X[0])')
#     plt.grid(True)
#     plt.legend()
#     plt.show()
#
# # 测试封装后的算法
# if __name__ == "__main__":
#     # 目标函数
#     def func(x):
#         return 100.0 * (x[0] - x[1] ** 2.0) ** 2.0 + (1 - x[1]) ** 2.0
#
#     # 生成初始的三个点
#     X = np.array([
#         [50, 50],
#         [10, 10],
#         [20, 20]
#     ])
#
#     # 调用 Nelder-Mead 算法
#     points, history_f = nelder_mead(func, X)
#
#     # 打印结果
#     print("最终迭代结果：")
#     # print(points)
#     # print(points[0][0][0])
#     print(f"x = {points[0][0][0]}, y = {points[0][0][1]}, f(x, y) = {points[1]}")
#
#     # 绘制收敛曲线
#     plot_convergence(history_f)
